from wxauto import WeChat
import requests
import time
import logging
from datetime import datetime
import json
import os

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 配置参数
DEEPSEEK_API_URL = "https://api.siliconflow.cn/v1/chat/completions"
DEEPSEEK_API_KEY = 'sk-accfnwnuyrxhohujjbszxcffrmrmhavhenaoxthhkteylbti'
MODEL = 'deepseek-ai/DeepSeek-V3'
MAX_TOKEN = 2000
TEMPERATURE = 1.3
API_TIMEOUT = 10  # API超时时间
POLLING_INTERVAL = 0.5  # 轮询间隔
MAX_CACHE_SIZE = 100  # 最大缓存数量

# 存储上下文和缓存
CHAT_CONTEXTS = {}
RESPONSE_CACHE = {}

def get_ai_response(message, user_name, new_messages=None):
    """获取AI回复"""
    try:
        # 检查缓存
        cache_key = f"{user_name}:{message}"
        if cache_key in RESPONSE_CACHE:
            logger.info("使用缓存回复")
            return RESPONSE_CACHE[cache_key]

        # 添加上下文
        if user_name not in CHAT_CONTEXTS:
            CHAT_CONTEXTS[user_name] = []

        # 如果上一条消息是用户发的，且时间间隔小于3秒，合并到上一条消息
        if (CHAT_CONTEXTS[user_name] and 
            CHAT_CONTEXTS[user_name][-1]["role"] == "user" and 
            time.time() - CHAT_CONTEXTS[user_name][-1].get("timestamp", 0) < 3):
            CHAT_CONTEXTS[user_name][-1]["content"] += "\n" + message
        else:
            CHAT_CONTEXTS[user_name].append({
                "role": "user", 
                "content": message,
                "timestamp": time.time()
            })

        # 保持上下文长度不超过5条消息
        if len(CHAT_CONTEXTS[user_name]) > 5:
            CHAT_CONTEXTS[user_name] = CHAT_CONTEXTS[user_name][-5:]

        headers = {
            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
            "Content-Type": "application/json"
        }

        # 读取prompt文件内容
        with open('prompt.md', 'r', encoding='utf-8') as f:
            prompt_content = f.read()

        data = {
            "model": MODEL,
            "messages": [
                {"role": "system", "content": prompt_content},
                *CHAT_CONTEXTS[user_name]
            ],
            "max_tokens": MAX_TOKEN,
            "temperature": TEMPERATURE
        }

        # 使用更短的超时时间
        response = requests.post(DEEPSEEK_API_URL, headers=headers, json=data, timeout=API_TIMEOUT)
        response.raise_for_status()
        reply = response.json()['choices'][0]['message']['content']

        # 添加回复到上下文
        CHAT_CONTEXTS[user_name].append({
            "role": "assistant", 
            "content": reply,
            "timestamp": time.time()
        })

        # 缓存回复
        RESPONSE_CACHE[cache_key] = reply
        if len(RESPONSE_CACHE) > MAX_CACHE_SIZE:
            RESPONSE_CACHE.pop(next(iter(RESPONSE_CACHE)))

        return reply
    except Exception as e:
        logger.error(f"调用AI API失败: {str(e)}")
        return "抱歉，我现在有点不舒服，稍后再聊..."

def check_wechat_window(wx):
    """检查微信窗口是否正常"""
    try:
        wx.SwitchToChat()
        return True
    except Exception as e:
        logger.error(f"微信窗口检测失败: {str(e)}")
        return False

def send_message(chat, message):
    """优化消息发送函数"""
    try:
        # 一次性发送完整消息
        chat.SendMsg(message)
        logger.info(f"消息发送成功: {message[:50]}...")
    except Exception as e:
        logger.error(f"消息发送失败: {str(e)}")
        # 如果发送失败，尝试分段发送
        for item in message.split('。'):
            if item.strip():
                chat.SendMsg(item.strip())

def listen_and_reply(contact_name):
    """监听并回复指定联系人的消息"""
    max_retries = 3
    retry_count = 0
    wx = None
    last_message_time = 0
    message_buffer = []
    is_processing = False

    while retry_count < max_retries:
        try:
            wx = WeChat()
            logger.info(f"成功连接到微信，开始监听 {contact_name} 的消息...")

            if not check_wechat_window(wx):
                raise Exception("微信窗口检测异常")

            wx.ChatWith(contact_name)
            time.sleep(1)  # 减少等待时间
            wx.AddListenChat(who=contact_name)
            logger.info(f"成功添加监听对象: {contact_name}")

            retry_count = 0

            while True:
                try:
                    if not check_wechat_window(wx):
                        raise Exception("微信窗口异常断开")

                    msgs = wx.GetListenMessage()
                    if not msgs:
                        time.sleep(POLLING_INTERVAL)  # 使用配置的轮询间隔
                        continue
                    
                    for chat in msgs:
                        chat_name = chat.who
                        if chat_name == contact_name:
                            messages = msgs[chat]
                            if not messages:
                                continue
                                
                            latest_msg = messages[-1]
                            if latest_msg.type == 'friend':
                                current_time = time.time()
                                
                                # 优化消息处理逻辑
                                if not is_processing:
                                    is_processing = True
                                    reply = get_ai_response(latest_msg.content, chat_name)
                                    send_message(chat, reply)
                                    is_processing = False
                                
                    time.sleep(POLLING_INTERVAL)  # 使用配置的轮询间隔
                except Exception as e:
                    logger.error(f"消息监听异常: {str(e)}")
                    continue

        except Exception as e:
            retry_count += 1
            logger.error(f"发生错误 (第{retry_count}次尝试): {str(e)}")
            
            if retry_count < max_retries:
                wait_time = 5  # 减少重试等待时间
                logger.info(f"{wait_time}秒后进行第{retry_count + 1}次重试...")
                time.sleep(wait_time)
            else:
                logger.error("已达到最大重试次数，程序退出")
                break

if __name__ == '__main__':
    CONTACT_NAME = "文件传输助手PRO"
    listen_and_reply(CONTACT_NAME) 